'''
Angle Encoder Interface.

Author: Hongjie Fang.
Modified by: Jin Liu.
Date: 2025/06/16
'''

import time
import serial
import numpy as np
from typing import List, Union

from ezrobot.sensor.base import BaseSensor
from ezrobot.utils.general import accurate_sleep


def hex2dex(e_hex):
    return int(e_hex, 16)

def hex2bin(e_hex):
    return bin(int(e_hex, 16))

def dex2bin(e_dex):
    return bin(e_dex)

def crc16(hex_num):
    """
    CRC16 verification
    :param hex_num:
    :return:
    """
    crc = '0xffff'
    crc16 = '0xA001'
    test = hex_num.split(' ')

    crc = hex2dex(crc)
    crc16 = hex2dex(crc16)
    for i in test:
        temp = '0x' + i
        temp = hex2dex(temp)
        crc ^= temp
        for i in range(8):
            if dex2bin(crc)[-1] == '0':
                crc >>= 1
            elif dex2bin(crc)[-1] == '1':
                crc >>= 1
                crc ^= crc16

    crc = hex(crc)
    crc_H = crc[2:4]
    crc_L = crc[-2:]

    return crc, crc_H, crc_L


class AngleEncoder(BaseSensor):

    def __init__(
        self,
        port,
        index: int, #Union[List[int], int],
        baudrate = 1000000,
        sleep_gap = 0.002,
        **kwargs
    ):
        self.index = index
        # self.index = index if isinstance(index, list) else [index]
        # self.index_map = {id:i for i, id in enumerate(self.index)}
        self.sleep_gap = sleep_gap

        self._dev = serial.Serial(port, baudrate=baudrate, timeout=0.4, **kwargs)
        if not self._dev.is_open:
            raise RuntimeError('Fail to open the serial port, please check your settings again.')
        self._dev.flushInput()
        self._dev.flushOutput()

        self.calibrate()

    # def recv_data(self, **kwargs):
    #     self._dev.flushInput()

    #     for i in self.index:
    #         sendbytes = str(i).zfill(2) + " 03 00 41 00 03"
    #         crc, crc_H, crc_L = crc16(sendbytes)
    #         sendbytes = sendbytes + ' ' + crc_L + ' ' + crc_H
    #         sendbytes = bytes.fromhex(sendbytes)
    #         self._dev.write(sendbytes)
    #         accurate_sleep(self.sleep_gap)

    #     recv_bytes = self._dev.read(len(self.index) * 11)
    #     if self._dev.inWaiting() > 0:
    #         tmp = self._dev.read_all()
    #         recv_bytes += tmp

    #     count = 0
    #     remains = self.index.copy()
    #     ret = np.zeros(len(self.index)).astype(np.float32)

    #     for i in range(len(recv_bytes) - 10):
    #         if recv_bytes[i + 1] == 3 and recv_bytes[i + 2] == 6 and recv_bytes[i] in remains:
    #             curr_id = recv_bytes[i]
    #             angle = 360 * (recv_bytes[i + 3] * 256 + recv_bytes[i + 4]) / 4096
    #             if angle > 350: # 临时解决编码器零点附近的突变问题
    #                 angle = 0
    #             ret[self.index_map[curr_id]] = angle
    #             count += 1
    #             remains.remove(curr_id)
    #     if count != len(self.index):
    #         raise RuntimeError(f'Failure to receive all encoders, errors occurred in ID {remains}')

    #     # map angle to gripper distance(radius:8mm,单边夹爪移动距离等于齿轮弧长，所以理论数据为2)
    #     ret = [angle * np.pi  / 180 * 8 * 2  for angle in ret]
    #     return ret

    def recv_data(self, **kwargs):
        self._dev.flushInput()

        sendbytes = str(self.index).zfill(2) + " 03 00 41 00 03"
 
        crc, crc_H, crc_L = crc16(sendbytes)
        sendbytes = sendbytes + ' ' + crc_L + ' ' + crc_H
        sendbytes = bytes.fromhex(sendbytes)
        self._dev.write(sendbytes)
        accurate_sleep(self.sleep_gap)
        re = self._dev.read(11)
        if self._dev.inWaiting() > 0:
            se = self._dev.read_all()
            re += se
        
        receive = False
        ret = 0
        for b in range(len(re) - 10):
            if re[b + 1] == 3 and re[b + 2] == 6 and re[b] == self.index:
                angle = 360 * (re[b + 3] * 256 + re[b + 4]) / 4096
                ret = angle
                receive = True
                if ret > 350: # 临时解决编码器零点附近的突变问题
                    ret = 0
        if not receive:
            return None
        ret = ret * np.pi  / 180 * 8 * 2 + 6.8 # 额外加的6.8部分是因为人皮肤弹性和夹爪的domain gap
        return ret

    def calibrate(self, index: Union[List[int], int] = None):
        """
        Calibrate the angle encoder to set the current angle as zero.
        """
        if index is None:
            index = self.index
        elif isinstance(index, int):
            index = [index]

        self._dev.flushInput()

        # for i in index:
        #     if i not in self.index:
        #         print(f"Encoder index {i} is not initialized.")
        #     sendbytes = str(i).zfill(2) + " 06 00 52 00 01"
        #     crc, crc_H, crc_L = crc16(sendbytes)
        #     sendbytes = sendbytes + ' ' + crc_L + ' ' + crc_H 
        #     sendbytes = bytes.fromhex(sendbytes) 
        #     self._dev.write(sendbytes) 
        #     accurate_sleep(self.sleep_gap)
        sendbytes = str(self.index).zfill(2) + " 06 00 52 00 01"
        crc, crc_H, crc_L = crc16(sendbytes)
        sendbytes = sendbytes + ' ' + crc_L + ' ' + crc_H 
        sendbytes = bytes.fromhex(sendbytes) 
        self._dev.write(sendbytes) 
        accurate_sleep(self.sleep_gap)

        self._dev.flushOutput()
        self._dev.flushInput()
        print("Angle encoder calibration done. ===")



if __name__ == '__main__':
    from ezrobot.sensor.base import SensorManager

    enc = AngleEncoder(port='COM17', index=1)
    enc.calibrate()
    encoder = SensorManager(sensor=enc, frequency=30, buffer_size=10, streaming=True)
    try:
        while True:
            time.sleep(0.1)
            print(encoder.get_data())
    except Exception as e:
        print(f"e")
        print("Stopping encoder streaming...")
        encoder.stop_streaming()